"""
Created on 03-Aug-2009

@author: dexen
"""

import cPickle


class Population(object):
    """A class to manage partially evaluated individuals in memory.
    
    """
    def __init__(self, inds_db, ind_db_lock, tasks_info):
        """Constructor
        
        Args:
        inds_db : A dict that stores all of the partially/fully evaluated
        individuals. Each individuals is made of dict as well. All the 
        individuals are pickled before stored and unpickled before retrieved.

        inds_db_lock : A shared lock for synchronizing access to the database.
    
        tasks_info : A dict that stores all the information about tasks.
        
        """
        self.inds_db = inds_db
        self.ind_db_lock = ind_db_lock
        self.tasks_info = tasks_info
        self.auto_id = 1 # WARNING: ! auto_id is supposed to be synchronized
    
    def get_ind(self, id):
        self.ind_db_lock.acquire()
        ind = cPickle.loads(self.inds_db[id])
        self.ind_db_lock.release()
        return ind
    
    def _get_task_obj(self, file_name, class_name):
        module = __import__(file_name, fromlist=[""])
        classobj = getattr(module, class_name)    
        return classobj()
    
    def _is_ind_for_task(self, ind, task):
        if not ind.__isalive__():
            return False 
        obj = task["obj"]
        return obj.__isinput__(ind)
    
    def _determine_queues(self, qdict, ind):
        """Assign the id of ind to the queues that can process.
        
        Args:
        qdict: A dict to store the list of ids for each task queue. 
        
        ind: An individual to be tested against all the task input format.
        
        Note:
        qdict will be returned by set_individuals to queue manager. Then the 
        queue manager will put the ids into queues accordingly.
        
        """
        for task_name, task in self.tasks_info.items():
            if self._is_ind_for_task(ind, task):
                qdict[task_name].append(ind.__id__)

    def _insert_ind(self, ind):
        ind.__id__ = self.auto_id #TODO: check if added to __modified__
        ind.revive() #TODO: check
        self.inds_db[self.auto_id] = cPickle.dumps(ind)
        self.auto_id += 1
        
    def _update_ind(self, ind):
        target_ind = cPickle.loads(self.inds_db[ind.__id__])
        for key in ind.__modified__.keys():
            setattr(target_ind, key, getattr(ind, key)) # TODO: check if works
        self.inds_db[ind.__id__] = cPickle.dumps(target_ind)
        return target_ind

    def get_individuals(self, task_name, ids):
        """Get individuals out of the inds_db database.
        
        task_name is used to select the corresponding fields of the 
        individuals.
            
        """
        res = []
        self.ind_db_lock.acquire()
        for id in ids:
            ind = cPickle.loads(self.inds_db[id])
            res.append(ind)
        self.ind_db_lock.release()
        return res

    def set_individuals(self, inds):
        """Put resulting individuals back to the inds_db.
        
        Returns a dict 'res' composed of (iqueue_name, list of ids) pairs.
        
        """
        res = dict.fromkeys(self.tasks_info, None)
        for key in res.keys(): 
            res[key] = []
        self.ind_db_lock.acquire()
        for ind in inds:
            if not ind._isnew():
                updated_ind = self._update_ind(ind)
                self._determine_queues(res, updated_ind)
            else: 
                self._insert_ind(ind)
                self._determine_queues(res, ind)
        self.ind_db_lock.release()
        return res
    
    def __repr__(self):
        res = ""
        for id, ind in self.inds_db.items():
            ind = cPickle.loads(ind)
            res += "%s: %s\n"%(id, ind)
        return res
